#include "tcp.h"
#include "tcp_sock.h"
#include "tcp_timer.h"

#include "log.h"
#include "ring_buffer.h"

#include <stdlib.h>
// update the snd_wnd of tcp_sock
//
// if the snd_wnd before updating is zero, notify tcp_sock_send (wait_send)


//exp3需要对`tcp_update_window`做出修改
/*
1. 记录更新前的tcp_tx_window_test结果
2. 更新snd_una， adv_wnd， snd_wnd。变量含义在tcp_sock结构体注释里。cwnd后面拥塞控制才会用到，可以设置一个较大的值0x7f7f7f7f。
3. 检查新的tcp_tx_window_test结果
4. 如果原本没有足够的窗口，现在有了，唤醒tsk->wait_send
*/



static inline void tcp_update_window(struct tcp_sock *tsk, struct tcp_cb *cb)
{
	int old_window_test = tcp_tx_window_test(tsk);

	//update snd_una adv_wnd snd_wnd cwnd
	tsk->snd_una = cb->ack;
	tsk->adv_wnd = cb->rwnd;
	
	log(DEBUG, "tcp_update_window: snd_una: %u adv_wnd: %u cwnd: %u", tsk->snd_una, tsk->adv_wnd, tsk->cwnd);
	tsk->snd_wnd = min((u32)tsk->cwnd, (u32)tsk->adv_wnd);

	if (less_than_32b(tsk->snd_wnd, (u32)TCP_MSS)) {
		log(DEBUG, "tcp_update_window: updated new snd_wnd: %u is less than TCP_MSS, set persist timer.",tsk->snd_wnd);
		tcp_set_persist_timer(tsk);
	}
	else{
		log(DEBUG," tcp_update_window: updated snd_wnd: %u is larger than TCP_MSS, cancel persist timer.", tsk->snd_wnd);
		tcp_unset_persist_timer(tsk);
	}

	int new_window_test = tcp_tx_window_test(tsk);

	if(old_window_test==0 &&new_window_test==1){//之前没有足够的窗口，现在有了，唤醒send
		log(DEBUG, "tcp_update_window: snd_wnd is updated, wake up wait_send.");
		wake_up(tsk->wait_send);
	}
	if(new_window_test==1){
		log(DEBUG,"tcp_update_window: snd_wnd is larger than TCP_MSS, wake up wait_send.");
		wake_up(tsk->wait_send);
	}
}

// update the snd_wnd safely: cb->ack should be between snd_una and snd_nxt
static inline void tcp_update_window_safe(struct tcp_sock *tsk, struct tcp_cb *cb)
{
	if (less_or_equal_32b(tsk->snd_una, cb->ack) && less_or_equal_32b(cb->ack, tsk->snd_nxt))
		tcp_update_window(tsk, cb);
	else{
		log(ERROR, "tcp_update_window_safe: cb->ack is not in the range of snd_una and snd_nxt.");
	}
}

#ifndef max
#	define max(x,y) ((x)>(y) ? (x) : (y))
#endif

// check whether the sequence number of the incoming packet is in the receiving
// window
static inline int is_tcp_seq_valid(struct tcp_sock *tsk, struct tcp_cb *cb)
{
	u32 rcv_end = tsk->rcv_nxt + max(tsk->rcv_wnd, 1);
	if (less_than_32b(cb->seq, rcv_end) && less_or_equal_32b(tsk->rcv_nxt, cb->seq_end)) {
		return 1;
	}
	else {
		log(ERROR, "received packet with invalid seq, drop it.");
		return 0;
	}
}




// Process the incoming packet according to TCP state machine. 
void tcp_process(struct tcp_sock *tsk, struct tcp_cb *cb, char *packet)
{

	//lock the tcp_process, use pthread_mutex_lock and unlock in the end.
	//log(DEBUG, "tcp_process start. lock tsk->sk_lock.");
	//pthread_mutex_lock(&tsk->sk_lock);

	//ADD: switcher: Routing to correct socket by cb info
	struct tcp_sock *tmp_tsk=NULL;

	if(cb){
	tmp_tsk = tcp_sock_lookup(cb);
	}

	if(!tmp_tsk){
		log(ERROR,"fail to find established or listen sock!");
	}
	else{//if found, routing to it
		//og(DEBUG, "find established or listen sock!");
		tsk = tmp_tsk;
	}

	if(cb&&tsk){
		char buf[64];
		tcp_copy_flags_to_str(cb->flags, buf);
		//log(DEBUG, "" IP_FMT ": %hu handle, coming packet flag is: %s", HOST_IP_FMT_STR(tsk->sk_sip), tsk->sk_sport, buf);
	}


	switch (tsk->state)
	{

	//起点 和 终点
	case TCP_CLOSED:
		if(!tsk){
			log(WARNING, "Tsk is not initialize, no need to release!");
		}

		log(DEBUG, "Connection closed, releasing resources");
		
		wake_up(tsk->wait_recv);
		
		//unhash, (free_sock is done in it)
		if(tsk){
			tcp_bind_unhash(tsk);
			tcp_unhash(tsk);
		}
		break;

	case TCP_LISTEN://server state
		//recv:SYN send: SYN, ACK
		if(cb->flags & TCP_SYN){
			//create new socket to manage connect
			struct tcp_sock *new_tsk = alloc_tcp_sock();
			if(!new_tsk){
				log(ERROR, "fail to alloc new sock!");
			}

			//set local info.
			new_tsk->local.ip = cb->daddr;
			new_tsk->local.port = tsk->local.port;
        	new_tsk->peer.ip = cb->saddr;
        	new_tsk->peer.port = cb->sport;
			new_tsk->parent = tsk;
	
			//enqueue the new accept tsk
			tcp_sock_accept_enqueue(new_tsk);
			
			//have something to accept, wake up server
			wake_up(tsk->wait_accept);

			//switch to state SYN_RECV
			tcp_set_state(new_tsk, TCP_SYN_RECV);

			// 4 elements set, we hash new_tsk(son of listen tsk) into established_table; preparing for routing
			tcp_hash(new_tsk);

			//sending message. 
			//work OK. packet can transmit from new_tsk and client
			//remember to update rcv_nxt, preparing for sending packet
			new_tsk->rcv_nxt = cb->seq + 1;
			tcp_send_control_packet(new_tsk, (TCP_SYN | TCP_ACK));
		}
		break;
	
	//RECV= received. AKA:RCVD
	case TCP_SYN_RECV:
		//for server, recv: ACK send: <nothing>
		if(cb->flags & TCP_ACK){
			// ACK recv
			// Update rcv_nxt
			tsk->rcv_nxt = cb->seq;
			
			// Update window
			tcp_update_window_safe(tsk, cb);

			tcp_set_state(tsk, TCP_ESTABLISHED);
			
			// 进入到 ESTABLISHED 取消重传计时器，唤醒等待连接的进程
			log(INFO, "tcp_process: TCP_ESTABLISHED, cancel retransmission timer.");
			tcp_unset_retrans_timer(tsk);

			//wake up recv in apps server
			wake_up(tsk->wait_recv);
		}
		break;
	
	case TCP_ESTABLISHED:
		
		if (cb->flags & TCP_ACK) {

			
			int removed=tcp_update_send_buffer(tsk, cb->ack);
			
			if(removed>0){
				log(INFO, "tcp_process: removed %d entries from send buffer.", removed);
			}
			else if(removed<0){
				log(ERROR, "tcp_process: remove entries from send buffer failed.");
			}
			else{
				log(DEBUG, "tcp_process: no entries removed from send buffer.");
			}

			if(!tsk->parent)
				tcp_congestion_control(tsk,cb,packet);
			
			// Update window
			tcp_update_window_safe(tsk, cb);

			

			// Check if the packet carries data or is a PSH
			// In this experiment 3, only server will recv data, client will not.
			if (cb->pl_len > 0 || (cb->flags & TCP_PSH)) {


				// 含有数据，将数据上送recv_ofo_buffer（乱序队列）
				//tsk->rcv_nxt=cb->seq;
				log(INFO,"tcp_process: recv seq: %u",cb->seq);
				int total_moved = tcp_recv_ofo_buffer_add_packet(tsk, cb);
				if (total_moved < 0) {
					log(ERROR, "tcp_process: failed to move data to receive buffer.");
					break;
				}
				log(INFO, "tcp_process: moved %d bytes to receive buffer.", total_moved);

			
				
				log(INFO,"Sending ack: %u, wnd is %u",tsk->rcv_nxt,tsk->rcv_wnd);
				tcp_send_control_packet(tsk, TCP_ACK); // Send ACK to peer

				// if (tsk->rcv_nxt == cb->seq) {
				// 	// Receive what is expected, write payload into rcv buf


				// 	int len = tcp_sock_write(tsk, cb->payload, cb->pl_len);
				// 	tsk->rcv_nxt = cb->seq+len;

				// 	// Send ACK to peer
				// 	tcp_send_control_packet(tsk, TCP_ACK);

				// 	// Wake up empty read
				// 	wake_up(tsk->wait_recv);
				// } else if (less_than_32b(tsk->rcv_nxt, cb->seq)) {
				// 	// Not ordered, deal later, now respond with ACK and discard data
				// 	log(DEBUG, "Not ordered, discard the data.");
				// 	tcp_send_control_packet(tsk, TCP_ACK);
				// }
			} 
			// 只有client收到ACK（数据传输阶段）才需要重设（这时候说明已经收到了回复ack，需要根据ack情况更新计时器）
			else if((cb->flags & TCP_ACK)&&cb->pl_len==0){
				log(DEBUG, "tcp_process: client recv ACK:%u, update retransmission timer.",cb->ack);

				// 可能需要更新窗口 唤醒一下read吧
				wake_up(tsk->wait_recv);

				tcp_update_retrans_timer(tsk);
			}
			else {
				// No data, handle ACK-only packet
				if (less_than_32b(cb->seq_end, tsk->rcv_nxt)) {
					// Receive ACKed segments, directly respond with ACK
					tcp_send_control_packet(tsk, TCP_ACK);
				}
			}
			
		}

		//for server: recv:FIN send:ACK to: CLOSE_WAIT
		if((cb->flags & TCP_FIN)){
			//compare with tsk->rcv_nxt, if less than it, drop it.
			if(cb->seq!=tsk->rcv_nxt){
				log(INFO, "tcp_process: not ordered FIN, drop it.");
				break;
			}

			tsk->rcv_nxt = cb->seq + 1; // Update rcv_nxt

			//if FIN|PSH
			if(cb->flags & TCP_PSH|| cb->pl_len > 0){
				//if payload is not empty, write payload into rcv buf
				int len=tcp_sock_write(tsk, cb->payload, cb->pl_len);
				if(len<0){
					log(ERROR, "tcp_process: write payload to rcv buf failed.");
					break;
				}
				tsk->rcv_nxt+=len;

				//send ACK to peer, note: we do not need to update rcv_nxt again.
				tcp_send_control_packet(tsk, TCP_ACK);

				//switch to CLOSE_WAIT state
				tcp_set_state(tsk, TCP_CLOSE_WAIT);
				//wake up empty read.
				wake_up(tsk->wait_recv);
			}
			else{//no data carry
				tcp_send_control_packet(tsk, TCP_ACK);
				//switch to CLOSE_WAIT state
				tcp_set_state(tsk, TCP_CLOSE_WAIT);
				wake_up(tsk->wait_recv);
			}
		}
		break;

	case TCP_FIN_WAIT_1:
		//For client: recv: ACK, send: <nothing> to: FIN_WAIT_2
		if(cb->flags & (TCP_ACK|TCP_FIN)){
			// Update send buf
			tsk->rcv_nxt =cb->seq+1; // Update rcv_nxt
			tcp_update_send_buffer(tsk, cb->ack);

			tcp_set_state(tsk,TCP_FIN_WAIT_2);
		}
		break;

	case TCP_FIN_WAIT_2:
		//For client: recv: FIN send:ACK to TIME_WAIT
		if(cb->flags & TCP_FIN){
			tsk->rcv_nxt =cb->seq+1; // Update rcv_nxt

			tcp_update_window_safe(tsk, cb);
			// 收到了应该的fin，取消重传计时器
			tcp_unset_retrans_timer(tsk);
			// 发送ack，进入time wait状态

			tcp_send_control_packet(tsk, TCP_ACK);
			tcp_set_state(tsk, TCP_TIME_WAIT);

		}
		break;

	case TCP_TIME_WAIT:
		//set timer, implement in tcp_timer
		log(DEBUG, "tcp_process: TCP_TIME_WAIT, set time-wait timer.");
		// 进入 TIME_WAIT 状态，设置时间等待计时器
		tcp_set_timewait_timer(tsk);

		break;
	
	//Another cycle: right side

	case TCP_SYN_SENT:
		//client, recv SYN, ACK, send ACK
		//for C client, Py server may remain in last core, and send you wrong ack. 
		//Py socket is always right. So we need to check the seq number due to it.
		if(cb->flags & (TCP_SYN | TCP_ACK )){
			// Update send buf
			tcp_update_send_buffer(tsk, cb->ack);

			// receive ACK, updating snd_window
			tcp_update_window_safe(tsk, cb);

			tsk->rcv_nxt = cb->seq + 1; // Update rcv_nxt

			// send ack to acknowledge server' s SYN
			tcp_send_control_packet(tsk, TCP_ACK);
			tsk->snd_nxt=tsk->snd_una;
			// 连接已建立，取消重传计时器
			tcp_unset_retrans_timer(tsk); 

			//wake up
			wake_up(tsk->wait_connect);
		}
		break;

	case TCP_CLOSE_WAIT:
		if(cb->flags & TCP_FIN){
			// Check if the sequence number is valid
			if((cb->seq!=tsk->rcv_nxt)){
				log(ERROR, "tcp_process: CLOSE_WAIT: cb->seq is not equal to tsk->rcv_nxt, drop it.");
				break;
			}
			//for client: recv: FIN send: ACK to: LAST_ACK
			tsk->rcv_nxt+=1;
			
			//if FIN|PSH
			if(cb->flags & TCP_PSH){
				
			}

		}
		break;

	case TCP_LAST_ACK:
		//for server: recv: ACK send: <nothing>
		if(cb->flags & TCP_ACK){
			// Update send buf
			tcp_update_send_buffer(tsk, cb->ack);

			tcp_set_state(tsk, TCP_CLOSED);

		}
		break;
	}

	//unlock the tcp_process, use pthread_mutex_unlock in the end.
	//pthread_mutex_unlock(&tsk->sk_lock);
	//log(DEBUG, "tcp_process end. unlock tsk->sk_lock.");
}



/*新增tcp_congestion_control函数
 函数tcp_congestion_control根据当前TCP拥塞控制的阶段(tsk->c_state)和收到的ACK数据包信息(cb->ack和ack_valid),
 更新拥塞窗口cwnd,慢启动阈值ssthresh等参数.
它通过状态机的方式处理不同的拥塞控制阶段: OPEN, DISORDER, LOSS 和 RECOVERY*/
void tcp_congestion_control(struct tcp_sock *tsk, struct tcp_cb *cb, char *packet)
{
	// Handle congestion control based on the current state
	switch (tsk->c_state) {
		case TCP_CONG_OPEN:
			// Congestion Avoidance Phase


			if (cb->ack > tsk->snd_una) {

				if(tsk->cwnd >= (float)tsk->ssthresh){
					// 进入拥塞避免阶段，线性增长
					
					tsk->cwnd += (float)((float)(TCP_MSS*TCP_MSS)/(float)tsk->cwnd); // Increase cwnd by MSS^2 / cwnd
				
				}
				else{
					// Slow Start Phase
					//log(INFO,"tcp_congestion_control: slow start. cwnd: %u -> %u",tsk->cwnd,tsk->cwnd+TCP_MSS);
					tsk->cwnd += (float)(1*TCP_MSS); // Increase cwnd by MSS
				}
	
			}

			if(cb->ack==tsk->snd_una){
				// Duplicate ACK received
				tsk->dupACKcount++;
				log(INFO, "TCP_CONG_OPEN: Duplicate ACK count: %d", tsk->dupACKcount);
				tcp_set_c_state(tsk, TCP_CONG_DISORDER); // Switch to DISORDER state
			}
			break;

		case TCP_CONG_DISORDER:
			//log(INFO,"TCP_CONG_DISORDER: start");
			// Fast Recovery Phase
			if(is_tcp_seq_valid(tsk, cb) == 0){
				log(ERROR, "tcp_congestion_control: cb->seq is not valid, drop it.");
				break;
			}

			if(cb->ack==tsk->snd_una){
				tsk->dupACKcount++;
				if (tsk->dupACKcount >= 3) {

					tsk->recovery_point=tsk->snd_nxt;

					tsk->ssthresh =tsk->cwnd / 2;
	
					tsk->cwnd = (float)(tsk->ssthresh + 3 * TCP_MSS);
					tcp_set_c_state(tsk, TCP_CONG_RECOVERY); // Switch to RECOVERY state
				}
	
			}
			else if(cb->ack>tsk->snd_una){//回到OPEN
				if(tsk->cwnd >= (float)tsk->ssthresh){
					// 进入拥塞避免阶段，线性增长
					
					tsk->cwnd += (float)((float)(TCP_MSS*TCP_MSS)/(float)tsk->cwnd); // Increase cwnd by MSS^2 / cwnd
				
				}
				else{
					// Slow Start Phase
					//log(INFO,"tcp_congestion_control: slow start. cwnd: %u -> %u",tsk->cwnd,tsk->cwnd+TCP_MSS);
					tsk->cwnd += (float)(1*TCP_MSS); // Increase cwnd by MSS
				}

				tsk->dupACKcount=0;
				tcp_set_c_state(tsk,TCP_CONG_OPEN);
			}
			break;

		case TCP_CONG_RECOVERY:
			// Slow Start Phase after Recovery 快恢复。（等待RTO很耗时间，直接重传）
			// 进行重传
			log(INFO, "TCP_CONG_RECOVERY: Retransmitting lost packets.");

			if (cb->ack > tsk->snd_nxt || cb->ack < tsk->snd_una) {
				log(WARNING, "Out-of-window ACK detected, ignoring.");
				return;
			}			

			// if new ACK, stop retrans, move to OPEN state
			if(cb->ack>tsk->snd_una){
				// 获得新的，跳出去
				tsk->cwnd = (float)(tsk->ssthresh);
				
				tsk->dupACKcount=0;
				tcp_set_c_state(tsk,TCP_CONG_OPEN);
			}
			if(tsk->dupACKcount>=3&&cb->ack == tsk->snd_una){
				tsk->cwnd+=(float)TCP_MSS;
			}

			tcp_retrans_send_buffer(tsk);
			//log(INFO,"tcp_congestion_control: Recovery, cwnd become half. %f -> %f",old_cwnd,tsk->cwnd);

			break;

		case TCP_CONG_LOSS:
			// Loss Recovery Phase,超时状态
			
			break;

		default:
			log(WARNING, "Unknown TCP congestion control state: %d", tsk->c_state);
			break;

	}

	static FILE *fp = NULL;
    static int time_us = 0;
    
    if (!fp) {
        fp = fopen("cwnd_by_ack.txt", "w");
        if (!fp) {
            log(ERROR, "Failed to open cwnd.txt");
            return;
        }
    }
	if (less_or_equal_32b(tsk->snd_una, cb->ack) && less_or_equal_32b(cb->ack, tsk->snd_nxt)){
         // 记录当前拥塞控制参数
         time_us += 1000;  // 假设每个ACK间隔1ms
         fprintf(fp, "%d %f %u %u %d\n", time_us, tsk->cwnd, tsk->ssthresh, tsk->adv_wnd,tsk->c_state);
         fflush(fp); 
        }
}